from dataclasses import dataclass
import os
from docx import Document
from pptx import Presentation
from typing import Callable, Dict, List
from MelodieFuncFlow import *
from .models import DocInfo, TableCell, Table
import hashlib
from .config import file_indexed


def get_md5(file_path: str) -> str:
    with open(file_path, "rb") as f:
        md5 = hashlib.md5()
        while True:
            data = f.read(8192)
            if not data:
                break
            md5.update(data)
    return md5.hexdigest()


class AlreadyIndexed(ValueError):
    pass


def extract_text_from_ppt(file_path):
    """
    提取 PowerPoint 文件中的所有文字内容，并返回 DocInfo 对象
    """
    # 创建 Presentation 对象

    md5_val = get_md5(file_path)
    if file_indexed(md5_val):
        raise AlreadyIndexed
    prs = Presentation(file_path)

    # 初始化一个列表存储表格单元格

    # 初始化一个列表存储段落文本
    paragraphs = []
    tables = []
    # 遍历每一个幻灯片
    for slide in prs.slides:
        # 遍历每一个形状
        for shape in slide.shapes:
            # 如果形状包含文字
            if hasattr(shape, "text"):
                # 将文字内容添加到段落列表中
                paragraphs.append(shape.text)

                # 处理表格
                if hasattr(shape, "table"):
                    table = shape.table
                    table_cells = []
                    # 遍历表格行和列
                    for row in range(len(table.rows)):
                        for col in range(len(table.columns)):
                            cell_value = table.cell(row, col).text
                            table_cells.append(TableCell(row, col, cell_value))
                    tables.append(Table(table_cells))
    # 计算文件 hash
    file_hash = hashlib.sha256()
    with open(file_path, "rb") as f:
        file_hash.update(f.read())
    file_hash = file_hash.hexdigest()

    # 获取文件创建和修改时间
    created_at = os.path.getctime(file_path)
    modified_at = os.path.getmtime(file_path)
    last_access = os.path.getatime(file_path)

    # 创建 DocInfo 对象
    doc_info = DocInfo(
        name=file_path,
        hash=file_hash,
        created_at=created_at,
        modified_at=modified_at,
        last_access=last_access,
        tables=(tables),
        paragraphs=(paragraphs),
    )

    return doc_info


def extract_text_from_docx(file_path) -> DocInfo:
    md5_val = get_md5(file_path)
    if file_indexed(md5_val):
        raise AlreadyIndexed
    doc = Document(file_path)
    paragraphs = []
    for paragraph in doc.paragraphs:
        paragraphs.append(paragraph.text)
    tables = []
    for table in doc.tables:
        cells = []
        for row in range(0, len(table.rows)):
            for col, cell in enumerate(table.row_cells(row)):  # 提取row行的全部列数据
                cells.append(TableCell(row, col, cell.text.strip()))
        tables.append(Table(cells))
    return DocInfo(
        file_path,
        md5_val,
        os.path.getctime(file_path),
        os.path.getmtime(file_path),
        os.path.getatime(file_path),
        tables,
        paragraphs,
    )


TEXT_EXTRACTORS: Dict[str, Callable[[str], DocInfo]] = {
    ".docx": extract_text_from_docx,
    ".doc": extract_text_from_docx,
    ".pptx": extract_text_from_ppt,
}

FILE_HANDLERS: Dict[str, Callable[[str], DocInfo]] = {
    ".docx": extract_text_from_docx,
    ".doc": extract_text_from_docx,
    ".pptx": extract_text_from_ppt,
}

SUPPORTED_EXTENSIONS = tuple(TEXT_EXTRACTORS.keys())

if __name__ == "__main__":
    file_path = r"C:\Users\houzh\Downloads\2023-12-30 北京市教委报告-0431.docx"
    text = extract_text_from_docx(file_path)
    print(text)
